package kafkax

import (
	"github.com/Shopify/sarama"
	"github.com/zeromicro/go-zero/core/logx"
	"sync"
)

type consumerGroupHandler struct {
	f         func(message *sarama.ConsumerMessage) error
	workerNum int
	stopChan  chan bool

	partition int32
	sync.Mutex

	waitSummit  []int64
	doneMessage map[int64]*sarama.ConsumerMessage
}

func (c *consumerGroupHandler) Setup(_ sarama.ConsumerGroupSession) error {
	logx.Info("consumerGroupHandler Setup")
	return nil
}
func (c *consumerGroupHandler) Cleanup(_ sarama.ConsumerGroupSession) error {
	logx.Info("consumerGroupHandler Cleanup")
	return nil
}

func (c *consumerGroupHandler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	workChan := make(chan *sarama.ConsumerMessage, c.workerNum)
	doneChan := make(chan *sarama.ConsumerMessage, c.workerNum)
	errChan := make(chan error, c.workerNum)

	c.waitSummit = make([]int64, 0)
	c.doneMessage = make(map[int64]*sarama.ConsumerMessage)

	logx.Info("ConsumeClaim start")

	wg := sync.WaitGroup{}
	go func() {
		wg.Add(1)
		defer wg.Done()
		// 清空workChan 关闭doneChan,errChan
		c.doWorker(workChan, doneChan, errChan)
	}()

	var err error
	go func() {
		wg.Add(1)
		defer wg.Done()
		// 关闭workChan
		err = c.doDistribute(claim, workChan, errChan)
	}()

	// 已完成的按照顺序 标记提交
	go func() {
		wg.Add(1)
		defer wg.Done()
		// 清空doneChan
		c.doSummit(sess, doneChan)
	}()
	wg.Wait()

	// 清空errChan
	for range errChan {
	}

	logx.Info("ConsumeClaim end", err)
	return err
}

func (c *consumerGroupHandler) doDistribute(
	claim sarama.ConsumerGroupClaim,
	workChan chan<- *sarama.ConsumerMessage,
	errChan <-chan error,
) error {
	defer close(workChan)

	for {
		select {
		case msg, ok := <-claim.Messages():
			if !ok {
				return nil
			}
			if msg == nil {
				continue
			}

			c.Lock()
			c.waitSummit = append(c.waitSummit, msg.Offset)
			c.Unlock()

			workChan <- msg
			logx.Info("doDistribute", msg.Partition, msg.Offset)
		case err := <-errChan:
			return err
		case <-c.stopChan:
			return nil
		}
	}
}

func (c *consumerGroupHandler) doSummit(
	sess sarama.ConsumerGroupSession,
	doneChan <-chan *sarama.ConsumerMessage,
) {
	for {
		select {
		case msg, ok := <-doneChan:
			if !ok {
				return
			}
			if msg == nil {
				continue
			}

			c.Lock()
			c.doneMessage[msg.Offset] = msg

			isSummit := false
			flag := true
			for flag && len(c.waitSummit) > 0 {
				flag = false
				for _, item := range c.doneMessage {
					if c.waitSummit[0] == item.Offset {
						logx.Info("doMark:", item.Partition, item.Offset)
						sess.MarkMessage(msg, "")
						if len(c.waitSummit) >= 1 {
							c.waitSummit = c.waitSummit[1:]
						}
						delete(c.doneMessage, item.Offset)
						isSummit = true
						flag = true
						break
					}
				}
			}
			if len(c.doneMessage) == 0 {
				c.doneMessage = make(map[int64]*sarama.ConsumerMessage)
			}
			c.Unlock()

			if isSummit {
				logx.Info("doSummit")
				sess.Commit()
			}
		}
	}
}

// 007的工作者
func (c *consumerGroupHandler) doWorker(
	workChan <-chan *sarama.ConsumerMessage,
	doneChan chan<- *sarama.ConsumerMessage,
	errChannel chan<- error,
) {
	defer close(doneChan)
	defer close(errChannel)

	sg := sync.WaitGroup{}
	for i := 0; i < c.workerNum; i++ {
		sg.Add(1)
		go func() {
			defer sg.Done()
			for {
				select {
				case msg, ok := <-workChan:
					if !ok {
						return
					}

					if msg == nil {
						continue
					}

					logx.Info("doWorker:", msg.Partition, msg.Offset)
					err := c.f(msg)
					if err == nil {
						doneChan <- msg
					} else {
						errChannel <- err
					}
				}
			}
		}()
	}
	sg.Wait()
}
